import os
from pathlib import Path

import numpy
import torch
from sklearn.model_selection import train_test_split

from rnn.torch.simple import TorchTest_simple
from rnn.torch.simple import TorchTrain_simple
from rnn.torch.simple.TorchRnnModule_simple import RnnModule


def run():
    r"""
    真实的数据关系

    .. math::

        y_0 = (x_{00}+x_{01}+...+x_{0d})/d

        y_t = (x_{t0}+x_{t1}+...+x_{td})/d

        y_t = y_{t-1}*0.5 + y_t*0.5

        y_t = y_t*3//1

    X:

    [[x_00,x_01,...,x_0d],

    [x_10,x_11,...,x_1d],

    ...

    [x_n0,x_n1,...,x_nd]]

    Y:

    [y_0,y_1,...,y_n]

    """
    x_dimension = 2
    x_list_maxlen = 5
    X = numpy.zeros((0, x_list_maxlen, x_dimension), dtype=numpy.float)
    Y = numpy.zeros((0, x_list_maxlen), dtype=numpy.int)
    for i in range(10000):
        x_list_len = numpy.random.randint(max(x_list_maxlen - 2, 1), x_list_maxlen + 1)
        x = numpy.random.rand(x_list_len, x_dimension)

        y = numpy.zeros((0,))
        for x_item in x:
            x_item_mean = numpy.mean(x_item)
            if y.shape[0] > 0:
                x_item_mean = y[y.shape[0] - 1] * 0.5 + x_item_mean * 0.5
            y = numpy.concatenate((y, [x_item_mean]))
        for index, y_item in enumerate(y):
            if y_item < 1.0 / 3:
                y[index] = 0
            elif y_item < 2.0 / 3:
                y[index] = 1
            else:
                y[index] = 2

        for j in range(x_list_maxlen - x_list_len):
            x = numpy.concatenate((x, numpy.zeros((1, x_dimension))))
            y = numpy.concatenate((y, numpy.zeros((1,)) - 1))

        y = y.astype(dtype=numpy.int)

        X = numpy.concatenate((X, [x]))
        Y = numpy.concatenate((Y, [y]))

    x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size=0.33, random_state=None)

    save_path = os.path.abspath('') + '/abc'
    if Path(save_path).is_file():
        model = torch.load(save_path)
    else:
        model = RnnModule(input_size=x_dimension)
        TorchTrain_simple.train(model, torch.tensor(x_train, dtype=torch.float),
                                torch.tensor(y_train, dtype=torch.long),
                                save_path)

    for parameter in model.parameters():
        print(parameter)

    model.eval()

    TorchTest_simple.test(model,
                          torch.tensor(x_test, dtype=torch.float),
                          torch.tensor(y_test, dtype=torch.long))


run()
